/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2019-2022.
 * Description: support lzma_improved decompression for squashfs
 * Author: yejianhua <yejianhua5@huawei.com>
 * Create: 2019-01-11
 */

#include <asm/unaligned.h>
#include <linux/slab.h>
#include <linux/vmalloc.h>

#include "lzma_improved.h"
#include "LzmaDec.h"

static void *size_alloc(ISzAllocPtr p, size_t size)
{
	if (size == 0)
		return 0;
	return vmalloc(size);
}

static void size_free(ISzAllocPtr p, const void *address)
{
	vfree(address);
}

static const ISzAlloc lzma_alloc = { size_alloc, size_free };

static int decode_improved(CLzmaDec *state, const unsigned char *input, int length,
	unsigned char *output, uint64_t unpack_size)
{
	int res = 0;
	size_t in_pos = 0;
	size_t out_pos = 0;
	size_t in_processed;
	size_t out_processed;
	bool there_is_size = (unpack_size != (uint64_t)(int64_t)-1);
	ELzmaFinishMode finish_mode;
	ELzmaStatus status;
	int length_remaining = length;
	int unpack_size_remaining = unpack_size;

	LzmaDec_Init(state);

	for (;;) {
		in_processed = min_t(size_t, IN_BUF_SIZE, length_remaining);
		out_processed = OUT_BUF_SIZE;
		finish_mode = LZMA_FINISH_ANY;

		if (there_is_size && out_processed > unpack_size_remaining) {
			out_processed = (size_t)unpack_size_remaining;
			finish_mode = LZMA_FINISH_END;
		}

		res = LzmaDec_DecodeToBuf(state, output + out_pos, &out_processed,
			input + in_pos, &in_processed, finish_mode, &status);

		in_pos += in_processed;
		out_pos += out_processed;
		length_remaining -= in_processed;
		unpack_size_remaining -= out_processed;

		if (res != SZ_OK || (there_is_size && unpack_size_remaining == 0))
			goto failed;

		if (in_processed == 0 && out_processed == 0) {
			if (there_is_size || status != LZMA_STATUS_FINISHED_WITH_MARK)
				res = SZ_ERROR_DATA;
			goto failed;
		}
	}

failed:
	return res;
}

int unlzma_improved(const unsigned char *input, int length, unsigned char *output)
{
	int i;
	int res = 0;
	uint64_t unpack_size;
	CLzmaDec state;

	/* header: 5 bytes of LZMA properties and 8 bytes of uncompressed size */
	unsigned char header[LZMA_PROPS_SIZE + 8];

	if (input == NULL || output == NULL)
		return SZ_ERROR_FAIL;

	if (length < (LZMA_PROPS_SIZE + 8))
		return SZ_ERROR_UNSUPPORTED;

	/* Read and parse header */
	memcpy(header, input, sizeof(header));

	unpack_size = 0;
	for (i = 0; i < 8; i++)
		unpack_size += (uint64_t)header[LZMA_PROPS_SIZE + i] << (i * 8);

	LzmaDec_Construct(&state);
	RINOK(LzmaDec_Allocate(&state, header, LZMA_PROPS_SIZE, &lzma_alloc));
	res = decode_improved(&state, input + sizeof(header), length - sizeof(header),
		output, unpack_size);
	LzmaDec_Free(&state, &lzma_alloc);

	return res;
}
